import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from scipy.spatial import Delaunay
from shapely.geometry import Polygon, MultiPolygon, MultiPoint, MultiLineString
from shapely.ops import polygonize, unary_union
import argparse


def parse_args():
    parser = argparse.ArgumentParser(description='t-SNE Visualization with Alpha Shape boundaries')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument('--arch', type=str, default='vit')
    parser.add_argument('--ft', type=str, default='full')
    parser.add_argument('--mode', type=str, default='test')
    parser.add_argument('--metrics', type=str, default='test')
    parser.add_argument('--metric-options', type=str, default='test')
    return parser.parse_args()


def alpha_shape(points, alpha):
    """
    Compute the alpha shape (concave hull) of a set of points.
    :param points: np.array of shape (n, 2) points.
    :param alpha: alpha value.
    :return: shapely.geometry.Polygon or None
    """
    if len(points) < 4:
        # Fallback to convex hull if less than 4 points
        return MultiPoint(list(points)).convex_hull

    tri = Delaunay(points)
    edges = set()

    # Loop over triangles
    for ia, ib, ic in tri.simplices:
        pa, pb, pc = points[ia], points[ib], points[ic]

        # Calculate lengths of sides of triangle
        a = np.linalg.norm(pa - pb)
        b = np.linalg.norm(pb - pc)
        c = np.linalg.norm(pc - pa)

        # Semiperimeter and area of the triangle
        s = (a + b + c) / 2.0
        area = max(s * (s - a) * (s - b) * (s - c), 1e-10)  # Avoid zero division

        # Circumcircle radius
        circ_radius = a * b * c / (4.0 * np.sqrt(area))
        if circ_radius < 1.0 / alpha:
            edges.add((ia, ib))
            edges.add((ib, ic))
            edges.add((ic, ia))

    # Construct the polygon(s)
    edge_points = [(points[edge[0]], points[edge[1]]) for edge in edges]
    m = MultiLineString(edge_points)
    triangles = list(polygonize(m))
    shape = unary_union(triangles)

    # Ensure a single Polygon
    if isinstance(shape, MultiPolygon):
        # Ensure MultiPolygon contains valid polygons
        polygons = list(shape.geoms)
        if polygons:
            # Choose the largest polygon
            shape = max(polygons, key=lambda p: p.area)
        else:
            print("MultiPolygon contains no valid polygons.")
            shape = None
    elif not isinstance(shape, Polygon):
        # Handle other geometry types (e.g., GeometryCollection)
        print("Generated shape is not a valid Polygon. Using convex hull as fallback.")
        shape = MultiPoint(list(points)).convex_hull

    return shape


def plot_alpha_shape(points, alpha, color):
    """
    Plot points and their alpha shape boundary.
    :param points: np.array of shape (n, 2) points.
    :param alpha: Alpha value for the alpha shape.
    :param color: Color of the boundary.
    """
    shape = alpha_shape(points, alpha)

    if isinstance(shape, Polygon):
        x, y = shape.exterior.xy
        plt.plot(x, y, linestyle="--", color=color, linewidth=2, alpha=0.8)
    else:
        print("No valid polygon generated.")

def main():
    args = parse_args()

    # Load embeddings and labels
    cls_embeddings = np.load(f'plots/tsne/ISIC/{args.ft}/ISIC_cls_embeddings_{args.arch}.npy')
    labels = np.load(f'plots/tsne/ISIC/{args.ft}/ISIC_labels_{args.arch}.npy')
    # cls_embeddings = np.load(f'plots/tsne/ISIC/{args.ft}/ISIC_cls_embeddings_o{args.arch}.npy')
    # labels = np.load(f'plots/tsne/ISIC/{args.ft}/ISIC_labels_o{args.arch}.npy')


    tsne = TSNE(n_components=2, perplexity=50, random_state=30)
    embeddings_2d = tsne.fit_transform(cls_embeddings)


    np.save(f'plots/tsne/ISIC/{args.ft}/ISIC_embeddings_2d_{args.arch}.npy', embeddings_2d)
    
    
    embeddings_2d = np.load(f'plots/tsne/ISIC/{args.ft}/ISIC_embeddings_2d_{args.arch}.npy')
    # Plot t-SNE results
    plt.figure(figsize=(10, 8))
    colors = ['g', 'b', 'r', 'y', 'm', 'c', 'k', 'orange', 'purple', 'pink', 'brown', 'gray', 'navy', 'teal', 'olive']
    unique_classes = np.unique(labels)
    # unique_classes = [2]
    scatter_handles = []
    for class_index in unique_classes:
        # Get embeddings for the current class
        class_mask = labels[:, 0] == class_index
        class_embeddings = embeddings_2d[class_mask]

        # Plot all points for the current class
        scatter = plt.scatter(class_embeddings[:, 0], class_embeddings[:, 1],color=colors[class_index], s=1, marker='o', label=f'Class {class_index}')
        scatter_handles.append(scatter)
        # Compute and plot Alpha Shape boundary
        plot_alpha_shape(class_embeddings, alpha=0.5, color=colors[class_index])

    # Remove borders and adjust layout
    plt.axis('off')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    legend = plt.legend(handles=scatter_handles, loc="upper right", fontsize=10)
    for handle in legend.legendHandles:
        handle.set_sizes([20])
    # Save the figure without borders
    
    plt.savefig(f'plots/tsne/ISIC/{args.arch}_{args.ft}.png', dpi=300, bbox_inches='tight', pad_inches=0)
    # plt.show()


if __name__ == '__main__':
    main()
